"""
邢不行 | 2023分享会
author: 邢不行
微信: xbx6660
"""
import pandas as pd
import plotly.graph_objs as go
from plotly.offline import plot
from plotly.subplots import make_subplots
import os

_ = os.path.abspath(os.path.dirname(__file__))  # 返回当前文件路径
root_path = os.path.abspath(os.path.join(_, '../'))  # 返回根目录文件夹


def get_trade_info(symbol_df, direction):
    """
    获取每一笔的交易信息
    :param symbol_df:交易数据
    :param direction: 交易方向
    :return:
    """
    # 先将数据拷贝一份避免污染数据源
    df = symbol_df.copy()

    # 标记周期开始时间
    df.loc[df['open_signal'].notnull(), 'open_time'] = df['candle_begin_time']
    df['open_time'].fillna(method='ffill', inplace=True)

    # 针对周期开始时间进行分组
    groups = df.groupby('open_time')

    res_list = []  # 存放每个周期的分组结果
    # 遍历每一个分组
    for t, g in groups:
        # 获取这笔交易的结束的索引
        end_inx = g[g['close_signal'] == 0].index.min()
        # 截取数据
        g = g[g.index <= end_inx]

        # 计算这笔交易的信息
        res = pd.DataFrame()
        res.loc[0, 'open_time'] = t  # 交易开始时间
        res.loc[0, 'close_time'] = g['candle_begin_time'].iloc[-1]  # 交易结束时间
        res.loc[0, 'open_price'] = g['open'].iloc[0]  # 开仓价格
        res.loc[0, 'close_price'] = g['close'].iloc[-1]  # 平仓价格

        res_list.append(res)

    # 聚合每一笔交易
    trade_df = pd.concat(res_list, ignore_index=True)
    # 添加交易方向
    trade_df['direction'] = 'long' if direction == '多头' else 'short'
    # 计算交易收益
    if direction == '多头':
        trade_df['return'] = trade_df['close_price'] / trade_df['open_price'] - 1
    else:
        trade_df['return'] = 1 - trade_df['close_price'] / trade_df['open_price']

    # 将收益率转为为百分比格式
    trade_df['return'] = trade_df['return'].apply(lambda x: str(round(100 * x, 2)) + '%')
    return trade_df


def draw_hedge_signal_plotly(df, save_path, title, day_df, trade_df, pic_size=[1880, 1656]):
    """
    绘制k线图以及指标图

    :param df:              原始数据
    :param save_path:       保存图片的路径
    :param title:           图片标题
    :param day_df:           日线数据
    :param pic_size:        图片大小
    """

    time_data = df['candle_begin_time']
    # 构建画布左轴
    fig = make_subplots(rows=3, cols=1)

    # 绘制k线图
    fig.add_trace(go.Candlestick(
        x=time_data,
        open=df['open'],  # 字段数据必须是元组、列表、numpy数组、或者pandas的Series数据
        high=df['high'],
        low=df['low'],
        close=df['close'],
        name='k线',
    ), row=1, col=1)

    # 绘制买卖点
    mark_point_list = []
    for i in df[(df['open_signal'].notna()) | (df['close_signal'].notna())].index:
        # 获取买卖点信息
        open_signal = df.loc[i, 'open_signal']
        close_signal = df.loc[i, 'close_signal']
        # 只有开仓信号，没有平仓信号
        if pd.notnull(open_signal) and pd.isnull(close_signal):
            signal = str(int(open_signal))
        # 没有开仓信号，只有平仓信号
        elif pd.isnull(open_signal) and pd.notnull(close_signal):
            signal = str(int(close_signal))
        else:  # 同时有开仓信号和平仓信号
            signal = str(int(open_signal)) + '_' + str(int(close_signal))
        # 标记买卖点，在最高价上方标记
        y = df.at[i, 'high']
        mark_point_list.append({
            'x': df.at[i, 'candle_begin_time'],
            'y': y,
            'showarrow': True,
            'text': signal,
            'arrowside': 'end',
            'arrowhead': 6,
        })

    # 绘制成交额
    fig.add_trace(go.Bar(x=time_data, y=df['quote_volume'], name='volume'), row=2, col=1)

    # 绘制日线
    time_data = day_df['candle_begin_time']
    # 绘制k线图
    fig.add_trace(go.Candlestick(
        x=time_data,
        open=day_df['open'],  # 字段数据必须是元组、列表、numpy数组、或者pandas的Series数据
        high=day_df['high'],
        low=day_df['low'],
        close=day_df['close'],
        name='日频k线',
    ), row=3, col=1)

    # mark_point_list3 = []
    for i in day_df[day_df['signal'].notna()].index:
        # 标记开始时间
        y = day_df.at[i, 'high']
        # 在周期数据上标记
        mark_point_list.append({
            'x': day_df.at[i, 'candle_begin_time'],
            'y': y,
            'showarrow': True,
            'text': day_df.loc[i, 'signal'],
            'arrowside': 'end',
            'arrowhead': 6,
        })
        # 在日线数据上标记
        mark_point_list.append({
            'x': day_df.at[i, 'candle_begin_time'],
            'y': y,
            'showarrow': True,
            'text': day_df.loc[i, 'signal'],
            'arrowside': 'end',
            'arrowhead': 6,
            'xref': 'x3', 'yref': 'y3'
        })

    table_trace = go.Table(header=dict(values=list(['open_time', 'close_time', 'return'])),
                           cells=dict(
                               values=[trade_df['open_time'], trade_df['close_time'], trade_df['return']]),
                           domain=dict(x=[0.73, 1.0], y=[0.57, 1.0]))
    fig.add_trace(table_trace)
    # 更新画布布局
    fig.update_layout(template="none", width=pic_size[0], height=pic_size[1], title_text=title, hovermode='x',
                      yaxis=dict(domain=[0.57, 1.0]), yaxis2=dict(domain=[0.43, 0.55]),
                      yaxis3=dict(domain=[0.0, 0.41]), xaxis=dict(domain=[0.0, 0.73]), xaxis2=dict(domain=[0.0, 0.73]),
                      xaxis3=dict(domain=[0.0, 0.73]), xaxis_rangeslider_visible=False,
                      xaxis3_rangeslider_visible=False)
    fig.update_layout(annotations=mark_point_list, title=title)

    # 保存路径
    save_path = save_path + title + '.html'
    plot(figure_or_data=fig, filename=save_path, auto_open=False)


def load_curve(bench, path, pct_dict, left_dict, right_dict, max_draw_down=False):
    """
    加载资金曲线
    :param bench:基准数据
    :param path:需要加载的资金曲线的路径
    :param pct_dict:{'本周期多空涨跌幅': 'ILLQ'}原始数据需要怎么重命名，主要是避免冲突
    :param left_dict:左轴的数据
    :param right_dict:右轴的数据
    :param max_draw_down:是否要计算最大回撤
    :return:
    """
    # 读取资金曲线
    df = pd.read_csv(path, encoding='gbk', parse_dates=['candle_begin_time'])
    # 对涨跌幅数据重命名，避免重名
    df.rename(columns=pct_dict, inplace=True)
    # 获取最新的涨跌幅列的名称
    new_col = list(pct_dict.values())[0]

    # 合并涨跌幅数据到基准数据上
    bench = pd.merge(bench, df[['candle_begin_time', new_col]], 'left', on='candle_begin_time')
    # 涨跌幅数据用0填充空值
    bench[new_col].fillna(value=0, inplace=True)
    # 计算资金曲线
    bench['净值_' + new_col] = (bench[new_col] + 1).cumprod()
    # 在左轴上加入需要绘制的信息
    left_dict['净值_' + new_col] = '净值_' + new_col

    # 如果需要绘制最大回撤
    if max_draw_down:
        # 计算历史至今的最大净值
        bench['m2h_' + new_col] = bench['净值_' + new_col].expanding().max()
        # 计算当前净值的回撤幅度
        bench['mdd_' + new_col] = bench['净值_' + new_col] / bench['m2h_' + new_col] - 1
        # 将回撤数据添加到右轴
        right_dict['最大回撤_' + new_col] = 'mdd_' + new_col
        # 删除必要的数据
        bench.drop(columns=['m2h_' + new_col], inplace=True)

        # 返回结果
    return bench, left_dict, right_dict


def draw_curves_plotly(df, data_dict, date_col=None, right_axis=None, pic_size=[1500, 800], chg=False,
                       title=None, path=root_path + '/data/pic.html', show=True):
    """
    绘制策略曲线
    :param df: 包含净值数据的df
    :param data_dict: 要展示的数据字典格式：｛图片上显示的名字:df中的列名｝
    :param date_col: 时间列的名字，如果为None将用索引作为时间列
    :param right_axis: 右轴数据 ｛图片上显示的名字:df中的列名｝
    :param pic_size: 图片的尺寸
    :param chg: datadict中的数据是否为涨跌幅，True表示涨跌幅，False表示净值
    :param title: 标题
    :param path: 图片路径
    :param show: 是否打开图片
    :return:
    """
    draw_df = df.copy()

    # 设置时间序列
    if date_col:
        time_data = draw_df[date_col]
    else:
        time_data = draw_df.index

    # 绘制左轴数据
    fig = make_subplots(specs=[[{"secondary_y": True}]])
    for key in data_dict:
        if chg:
            draw_df[data_dict[key]] = (draw_df[data_dict[key]] + 1).fillna(1).cumprod()
        fig.add_trace(go.Scatter(x=time_data, y=draw_df[data_dict[key]], name=key, ))

    # 绘制右轴数据
    if right_axis:
        for key in list(right_axis.keys()):
            fig.add_trace(go.Scatter(x=time_data, y=draw_df[right_axis[key]], name=key + '(右轴)',
                                     opacity=0.5, line=dict(width=0), fill='tozeroy',
                                     yaxis='y2'))  # 标明设置一个不同于trace1的一个坐标轴
    fig.update_layout(template="none", width=pic_size[0], height=pic_size[1], title_text=title,
                      hovermode="x unified", hoverlabel=dict(bgcolor='rgba(255,255,255,0.5)', ),
                      )
    fig.update_layout(
        updatemenus=[
            dict(
                buttons=[
                    dict(label="线性 y轴",
                         method="relayout",
                         args=[{"yaxis.type": "linear"}]),
                    dict(label="Log y轴",
                         method="relayout",
                         args=[{"yaxis.type": "log"}]),
                ])], )
    plot(figure_or_data=fig, filename=path, auto_open=False)

    fig.update_yaxes(
        showspikes=True, spikemode='across', spikesnap='cursor', spikedash='solid', spikethickness=1,  # 峰线
    )
    fig.update_xaxes(
        showspikes=True, spikemode='across+marker', spikesnap='cursor', spikedash='solid', spikethickness=1,  # 峰线
    )

    # 打开图片的html文件，需要判断系统的类型
    if show:
        res = os.system('start ' + path)
        if res != 0:
            os.system('open ' + path)
